#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "tinyndarray.h"

namespace py = pybind11;
namespace tn = tinyndarray;

// 将DtypeTag暴露给Python
void init_dtype_tag(py::module& m) {
    py::enum_<tn::DtypeTag>(m, "DtypeTag")
            .value("FP32", tn::FP32)
            .value("FP64", tn::FP64)
            .value("INT32", tn::INT32)
            .value("BOOL", tn::BOOL);
}

// 将Dtype暴露给Python
void init_dtype(py::module& m) {
    py::class_<tn::Dtype>(m, "Dtype")
            .def(py::init<tn::DtypeTag>())
            .def_readwrite("tag", &tn::Dtype::tag)
            .def("__eq__", [](const tn::Dtype& a, const tn::Dtype& b) { return a.tag == b.tag; })
            .def("__repr__", [](const tn::Dtype& d) {
                switch (d.tag) {
                    case tn::FP32: return "float32";
                    case tn::FP64: return "float64";
                    case tn::INT32: return "int32";
                    case tn::BOOL: return "bool";
                    default: return "unknown";
                }
            });
}

// 将NdArray暴露给Python
void init_ndarray(py::module& m) {
    py::class_<tn::NdArray>(m, "NdArray")
            // 构造函数
            .def(py::init<>())
            .def(py::init<tn::Dtype, const tn::Shape&>())
            .def(py::init<tn::Dtype, const tn::Shape&, float32_t>())

            // 属性
            .def_property_readonly("shape", static_cast<const tn::Shape& (tn::NdArray::*)() const>(&tn::NdArray::shape))
            .def_property_readonly("dtype", &tn::NdArray::dtype)
            .def_property_readonly("ndim", &tn::NdArray::ndim)
            .def_property_readonly("size", &tn::NdArray::size)
            .def_property_readonly("itemsize", &tn::NdArray::itemsize)
            .def_property_readonly("strides", &tn::NdArray::strides)
            .def_property_readonly("data_ptr",
                                   [](const tn::NdArray& self) { return reinterpret_cast<uintptr_t>(self.data_ptr()); })

            // 方法
            .def("copy", &tn::NdArray::copy)
            .def("astype", &tn::NdArray::astype)
            .def("reshape", &tn::NdArray::reshape)
            .def("flatten", &tn::NdArray::flatten)
            .def("ravel", &tn::NdArray::ravel)
            .def("to_string", &tn::NdArray::to_string)
            // .def("fill",
            //      [](tn::NdArray& self, py::object value) {
            //          if (PyFloat_Check(value.ptr())) {
            //              self.fill(value.cast<float>());
            //          } else if (PyLong_Check(value.ptr())) {
            //              self.fill(value.cast<int>());
            //          } else if (PyBool_Check(value.ptr())) {
            //              self.fill(value.cast<bool>());
            //          }
            //      })
            .def(
                    "sum",
                    [](const tn::NdArray& self, py::object axis, bool keepdims) {
                        if (axis.is_none()) {
                            return self.sum(tn::Axis(), keepdims);
                        } else if (py::isinstance<py::tuple>(axis)) {
                            tn::Axis ax;
                            for (auto a : axis.cast<py::tuple>()) {
                                ax.push_back(a.cast<int>());
                            }
                            return self.sum(ax, keepdims);
                        } else {
                            return self.sum(tn::Axis{axis.cast<int>()}, keepdims);
                        }
                    },
                    py::arg("axis") = py::none(), py::arg("keepdims") = false)
            // 索引访问 - 修改getitem返回实际值
            .def("__getitem__",
                 [](const tn::NdArray& self, size_t index) -> py::object {
                     auto proxy = self[index];
                     switch (self.dtype().tag) {
                         case tn::FP32: return py::cast(static_cast<float32_t>(proxy));
                         case tn::FP64: return py::cast(static_cast<float64_t>(proxy));
                         case tn::INT32: return py::cast(static_cast<int32_t>(proxy));
                         case tn::BOOL: return py::cast(static_cast<bool>(proxy));
                         default: throw std::runtime_error("Unsupported dtype");
                     }
                 })
            .def("__getitem__",
                 [](const tn::NdArray& self, py::tuple indices) -> py::object {
                     std::vector<int> idx;
                     for (auto item : indices) {
                         idx.push_back(item.cast<int>());
                     }
                     auto proxy = self[idx];
                     switch (self.dtype().tag) {
                         case tn::FP32: return py::cast(static_cast<float32_t>(proxy));
                         case tn::FP64: return py::cast(static_cast<float64_t>(proxy));
                         case tn::INT32: return py::cast(static_cast<int32_t>(proxy));
                         case tn::BOOL: return py::cast(static_cast<bool>(proxy));
                         default: throw std::runtime_error("Unsupported dtype");
                     }
                 })
            .def("__setitem__",
                 [](tn::NdArray& self, const std::vector<int>& idx, py::object value) {
                     // 检查 Python 原生类型
                     py::module ctypes = py::module::import("ctypes");
                     py::object c_float_type = ctypes.attr("c_float");
                     if (py::isinstance(value, c_float_type)) {
                         self[idx] = value.cast<float>();
                     } else if (py::isinstance<py::int_>(value)) {
                         self[idx] = value.cast<int>();
                     } else if (py::isinstance<py::bool_>(value)) {
                         self[idx] = value.cast<bool>();
                     }
                     // 检查 NumPy 标量类型
                     else if (py::isinstance<py::array_t<float>>(value)) {
                         auto arr = value.cast<py::array_t<float>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<double>>(value)) {
                         auto arr = value.cast<py::array_t<double>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<int>>(value)) {
                         auto arr = value.cast<py::array_t<int>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<bool>>(value)) {
                         auto arr = value.cast<py::array_t<bool>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     }
                     // 尝试通用转换
                     else {
                         throw std::runtime_error("Unsupported value type for assignment");
                     }
                 })
            .def("__setitem__",
                 [](tn::NdArray& self, size_t idx, py::object value) {
                    py::module ctypes = py::module::import("ctypes");
                     py::object c_float_type = ctypes.attr("c_float");
                     if (py::isinstance(value, c_float_type)) {
                         self[idx] = value.attr("value").cast<float>();
                     } else
                     if (py::isinstance<py::float_>(value)) {
                         self[idx] = value.cast<double>();
                     } else if (py::isinstance<py::int_>(value)) {
                         self[idx] = value.cast<int>();
                     } else if (py::isinstance<py::bool_>(value)) {
                         self[idx] = value.cast<bool>();
                     }
                     // 检查 NumPy 标量类型
                     else if (py::isinstance<py::array_t<float>>(value)) {
                         auto arr = value.cast<py::array_t<float>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<double>>(value)) {
                         auto arr = value.cast<py::array_t<double>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<int>>(value)) {
                         auto arr = value.cast<py::array_t<int>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     } else if (py::isinstance<py::array_t<bool>>(value)) {
                         auto arr = value.cast<py::array_t<bool>>();
                         if (arr.size() != 1) {
                             throw std::runtime_error("Expected a scalar value");
                         }
                         self[idx] = *arr.data();
                     }
                     // 尝试通用转换
                     else {
                         throw std::runtime_error("Unsupported value type for assignment");
                     }
                 })
            // 特殊方法（运算符）
            // 数组与数组运算
            .def("__add__", [](const tn::NdArray& a, const tn::NdArray& b) { return a + b; })
            .def("__sub__", [](const tn::NdArray& a, const tn::NdArray& b) { return a - b; })
            .def("__mul__", [](const tn::NdArray& a, const tn::NdArray& b) { return a * b; })
            .def("__truediv__", [](const tn::NdArray& a, const tn::NdArray& b) { return a / b; })
            .def("__eq__", [](const tn::NdArray& a, const tn::NdArray& b) { return a == b; })
            // .def("__ne__", [](const tn::NdArray& a, const tn::NdArray& b) { return a != b; })
            // .def("__lt__", [](const tn::NdArray& a, const tn::NdArray& b) { return a < b; })
            // .def("__le__", [](const tn::NdArray& a, const tn::NdArray& b) { return a <= b; })
            // .def("__gt__", [](const tn::NdArray& a, const tn::NdArray& b) { return a > b; })
            // .def("__ge__", [](const tn::NdArray& a, const tn::NdArray& b) { return a >= b; })
            .def("__neg__", [](const tn::NdArray& a) { return -a; })

            // 数组与标量运算（左操作数）
            .def("__add__", [](const tn::NdArray& a, float b) { return a + b; })
            // .def("__add__", [](const tn::NdArray& a, int b) { return a + b; })
            // .def("__sub__", [](const tn::NdArray& a, float b) { return a - b; })
            // .def("__sub__", [](const tn::NdArray& a, int b) { return a - b; })
            .def("__mul__", [](const tn::NdArray& a, float b) { return a * b; })
            // .def("__mul__", [](const tn::NdArray& a, int b) { return a * b; })
            .def("__truediv__", [](const tn::NdArray& a, float b) { return a / b; })
            // .def("__truediv__", [](const tn::NdArray& a, int b) { return a / b; })

            // // 标量与数组运算（右操作数）
            // .def("__radd__", [](const tn::NdArray& a, float b) { return b + a; })
            // .def("__radd__", [](const tn::NdArray& a, int b) { return b + a; })
            // .def("__rsub__", [](const tn::NdArray& a, float b) { return b - a; })
            // .def("__rsub__", [](const tn::NdArray& a, int b) { return b - a; })
            // .def("__rmul__", [](const tn::NdArray& a, float b) { return b * a; })
            // .def("__rmul__", [](const tn::NdArray& a, int b) { return b * a; })
            // .def("__rtruediv__", [](const tn::NdArray& a, float b) { return b / a; })
            // .def("__rtruediv__", [](const tn::NdArray& a, int b) { return b / a; })
            // 工厂方法
            .def_static("zeros", &tn::NdArray::Zeros)
            .def_static("ones", &tn::NdArray::Ones)
            .def_static("empty", &tn::NdArray::Empty);
}

// 主模块
PYBIND11_MODULE(_core, m) {
    m.doc() = "Core C++ implementation for tinyndarray";

    // 初始化各个子模块
    init_dtype_tag(m);
    init_dtype(m);
    init_ndarray(m);

    // 添加常量
    m.attr("float32") = tn::Dtype(tn::FP32);
    m.attr("float64") = tn::Dtype(tn::FP64);
    m.attr("int32") = tn::Dtype(tn::INT32);
    m.attr("bool_") = tn::Dtype(tn::BOOL);

    // 随机数
    m.def(
            "seed",
            [](py::object seed_obj) {
                if (seed_obj.is_none()) {
                    tn::NdArray::Seed();
                } else {
                    tn::NdArray::Seed(seed_obj.cast<uint32_t>());
                }
            },
            py::arg("seed") = py::none());

    m.def(
            "uniform",
            [](float low, float high, const tn::Shape& shape) { return tn::NdArray::Uniform(low, high, shape); },
            py::arg("low") = 0.0f, py::arg("high") = 1.0f, py::arg("shape"));

    // 文件I/O
    m.def(
            "save", [](const std::string& file_path, const tn::NdArray& arr) { arr.Save(file_path); },
            py::arg("file_path"), py::arg("array"));

    m.def("load", &tn::Load, py::arg("file_path"));
}